package com.yao2san.simdif.service.synchronizer;

import cn.hutool.core.util.BooleanUtil;
import cn.hutool.db.sql.Condition;
import cn.hutool.db.sql.SqlBuilder;
import com.yao2san.simdif.entity.ColumnInfo;
import com.yao2san.simdif.entity.DatasourceConnectionAttr;
import com.yao2san.simdif.entity.TableInfo;
import com.yao2san.simdif.enums.Drivers;
import com.yao2san.simdif.service.DatasourceConnectHelper;
import lombok.extern.slf4j.Slf4j;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.*;
import java.util.stream.Collectors;

/**
 * PgSql数据源同步
 *
 * @author wxg
 */
@Slf4j
public class PgsqlDatasourceSynchronizer extends AbstractDatasourceSynchronizer {

    private static final String SQL_QUERY_TABLE = "select t.*, " +
            "       cast(obj_description(relfilenode, 'pg_class') AS VARCHAR) AS \"TABLE_COMMENT\" " +
            "from information_schema.tables t, " +
            "     pg_class pc " +
            "where table_catalog=? " +
            "  and table_schema =? " +
            "  and t.table_name = pc.relname " +
            "  and t.table_type = 'BASE TABLE' " +
            "  AND t.table_name NOT LIKE 'pg_%' " +
            "  AND t.table_name NOT LIKE 'sql_%'" +
            "  and pc.relkind in('r','p') ";
    private static final String SQL_QUERY_COLUMN = "SELECT c.table_name                            as \"TABLE_NAME\", " +
            "       c.ordinal_position                      as \"ORDINAL_POSITION\", " +
            "       c.data_type                             as \"COLUMN_DATA_TYPE\", " +
            "       c.column_default                        as \"COLUMN_DEFAULT\", " +
            "       col_description(pa.attrelid, pa.attnum) as \"COLUMN_COMMENT\", " +
            "       format_type(pa.atttypid, pa.atttypmod)  as \"COLUMN_TYPE\", " +
            "       pa.attname                              as \"COLUMN_NAME\", " +
            "       pa.attnotnull                           as \"IS_NOT_NULL\" " +
            "FROM pg_class as pc, " +
            "     pg_attribute as pa, " +
            "     information_schema.columns c " +
            "where pa.attrelid = pc.oid " +
            "  and c.table_name = pc.relname " +
            "  and c.column_name = pa.attname " +
            "  and pa.attnum > 0 " +
            "  and pc.relkind in('r','p') " +
            "  AND pc.relname NOT LIKE 'pg_%' " +
            "  AND pc.relname NOT LIKE 'sql_%' " +
            "  and c.table_catalog=? " +
            "  and c.table_schema = ?";
    private static final String SQL_QUERY_INDEX = "SELECT n.nspname     AS SCHEMA, " +
            "       c.relname     AS TABLE_NAME, " +
            "       array_to_string(array_agg(i.relname || SUBSTRING(pg_get_indexdef(i.oid), position('(' in pg_get_indexdef(i.oid)))), '; ') as \"INDEX_NAME\" " +
            "FROM pg_index x " +
            "         JOIN pg_class c ON c.oid = x.indrelid " +
            "         JOIN pg_class i ON i.oid = x.indexrelid " +
            "         LEFT JOIN pg_namespace n ON n.oid = c.relnamespace " +
            "         LEFT JOIN pg_tablespace t ON t.oid = i.reltablespace " +
            "where n.nspname = ? " +
            "  and c.relkind in ('r', 'p') " +
            "group by SCHEMA, TABLE_NAME";


    private final static ThreadLocal<Map<String, List<ColumnInfo>>> COLUMN_INFOS = new ThreadLocal<>();

    public PgsqlDatasourceSynchronizer(String url, String user, String password) {
        super(url, user, password);
    }


    @Override
    public List<TableInfo> getTables(Connection connection, String database) {
        List<TableInfo> allTables = getAllTables(connection, database);
        log.info("Find {} tables", allTables.size());
        Map<String, List<String>> allTableIndexes = getAllTableIndexes(connection, database);
        log.info("Find {} indexes", allTableIndexes.size());
        for (TableInfo tableInfo : allTables) {
            List<String> indexes = allTableIndexes.get(tableInfo.getTableName());
            if (indexes != null) {
                tableInfo.setTableIndex(String.join("; ", indexes));
            }
        }
        return allTables;
    }

    private List<TableInfo> getAllTables(Connection connection, String database) {
        PreparedStatement statement;
        try {
            statement = connection.prepareStatement(SQL_QUERY_TABLE);
            DatasourceConnectionAttr connectionAttr = DatasourceConnectHelper.getConnectionAttr();
            statement.setString(1, connectionAttr.getCatalog());
            statement.setString(2, connectionAttr.getSchema());
            ResultSet resultSet = statement.executeQuery();
            List<TableInfo> tableInfos = new ArrayList<>();
            while (resultSet.next()) {
                TableInfo tableInfo = new TableInfo();
                tableInfo.setTableName(resultSet.getString("TABLE_NAME"));
                tableInfo.setTableType(resultSet.getString("TABLE_TYPE"));
                tableInfo.setTableComment(resultSet.getString("TABLE_COMMENT"));
                tableInfos.add(tableInfo);
            }
            return tableInfos;
        } catch (SQLException e) {
            throw new RuntimeException(e);
        }
    }

    private Map<String, List<String>> getAllTableIndexes(Connection connection, String database) {
        Map<String, List<String>> indexMap = new HashMap<>();
        PreparedStatement statement;
        try {
            statement = connection.prepareStatement(SQL_QUERY_INDEX);
            DatasourceConnectionAttr connectionAttr = DatasourceConnectHelper.getConnectionAttr();
            statement.setString(1, connectionAttr.getSchema());
            List<String> indexes = new ArrayList<>();
            ResultSet indexResultSet = statement.executeQuery();
            while (indexResultSet.next()) {
                String tableName = indexResultSet.getString("TABLE_NAME");
                String indexName = indexResultSet.getString("INDEX_NAME");
                List<String> list = indexMap.computeIfAbsent(tableName, k -> new ArrayList<>());
                list.add(indexName);
            }
            return indexMap;
        } catch (SQLException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public List<ColumnInfo> getColumns(Connection connection, String database, String tableName) {
        if (COLUMN_INFOS.get() == null) {
            cacheAllColumns(connection, database);
        }
        return COLUMN_INFOS.get().get(tableName);
    }

    private void cacheAllColumns(Connection connection, String database) {
        List<ColumnInfo> columnInfos = new ArrayList<>();
        PreparedStatement statement;
        try {
            statement = connection.prepareStatement(SQL_QUERY_COLUMN);
            DatasourceConnectionAttr connectionAttr = DatasourceConnectHelper.getConnectionAttr();
            statement.setString(1, connectionAttr.getCatalog());
            statement.setString(2, connectionAttr.getSchema());
            ResultSet resultSet = statement.executeQuery();
            while (resultSet.next()) {
                ColumnInfo columnInfo = new ColumnInfo();
                columnInfo.setColumnName(resultSet.getString("COLUMN_NAME"));
                columnInfo.setColumnSort(resultSet.getInt("ORDINAL_POSITION"));
                columnInfo.setColumnDefault(resultSet.getObject("COLUMN_DEFAULT"));
                columnInfo.setColumnType(resultSet.getString("COLUMN_TYPE"));
                columnInfo.setColumnDataType(resultSet.getString("COLUMN_DATA_TYPE"));
                columnInfo.setColumnComment(resultSet.getString("COLUMN_COMMENT"));
                columnInfo.setColumnNull(!BooleanUtil.toBoolean(resultSet.getString("IS_NOT_NULL")));
                columnInfo.setTableName(resultSet.getString("TABLE_NAME"));
                columnInfos.add(columnInfo);
            }
            //按表名分组
            Map<String, List<ColumnInfo>> columns = columnInfos.stream().collect(Collectors.groupingBy(ColumnInfo::getTableName));
            COLUMN_INFOS.set(columns);
        } catch (SQLException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public String getDriverFile() {
        return Drivers.PGSQL.getDriver();
    }

    @Override
    public String getDriverClassName() {
        return Drivers.PGSQL.getDriverClassName();
    }

    @Override
    public void close() throws SQLException {
        super.close();
        COLUMN_INFOS.remove();
    }
}
